'''
Created on Feb 2, 2010

@author: oabalbin
'''

#import sys;
#import string;
#import StringIO;
import os;
import sys;
import glob;
import MySQLdb;
import numpy as np
from optparse import OptionParser
import operator
from collections import deque

# My libraries
import signatures.parsers.read_microarray_data as pm
import signatures.parsers.read_gene_lists as gpm
import signatures.db.tables as dt
import signatures.db.query as dq
import signatures.common.classes as cm #record, arrayMatrix
import signatures.preprocess.tools as tl #filter_constant_genes, get_variable_genes, 
import signatures.seedgenes.seedlist as sd #get_seed_genes, write_varin_file
import signatures.bfrm.paramfile as bp
import signatures.bfrm.bfrm_run as br

class pipeline():
    
    def __init__(self, dbhost, dbuser, dbpasswd, database):
        
        self.mt = dt.tables(dbhost, dbuser, dbpasswd, database)
        self.mq = dq.query(dbhost, dbuser, dbpasswd, database)
        #table information
        self.tableDescription = 'probeID VARCHAR(40), geneSym VARCHAR(40), sample VARCHAR(40), log_fold FLOAT(20,14), INDEX probe (probeID), INDEX gene (geneSym), INDEX sample (sample)'
        #self.tableName = 'ERG_Prostate'
        #self.tableName = 'rnaseq_breast_cl_jun'
        #self.tableName = 'rnaseq_prostate_AllTUs'
        self.tableName = 'rnaseq_prostate_AllTUsv5'
        #self.tableName = 'erg_signatures_only'
        # Parameters for pre-processing data:
        # percentage of low variable gene to be filter out of the analysis
        self.lvar=0
        # percentage of low median gene to be filter out of the analysis
        self.lmed=0
        
    
    
    def parse_input_annotfiles(self, annotInputFile):
    
        seed_file, outdirs, samples, factors, typefile = {},{},{},{},'nofile'
    
        for line in annotInputFile:
            fields = line.strip('\n')
                 
            if fields[0][0] == '' or fields[0][0] == '#':
                continue
            
            if fields[0][0] == '@':
                typefile = fields
                continue
            
            fields = fields.split('=')
            dict = {'@Seed genes':seed_file,'@Samples':samples,'@Output Directories':outdirs,'@Factors':factors}[typefile]
            if typefile=='@Factors':
                dict[fields[0]] = map(int,fields[1].split(','))
            else:
                dict[fields[0]] = fields[1]
            
            
        return seed_file, samples, outdirs, factors
    
   
                
def get_expMat(pipe, samplelist, genelist, normal_samples,norm=False,row_aveg=False):
    """
    Returns a expMat, an expression Array object (common classes) object
    input: pipe= pipeline object, samplelist = a list of samples, genelist=list of genes
    norm = bool indicates if normalization of expression data or or just log2(expression)
    """
    
    pipe.mq.create_connector()
    expMat = pipe.mq.get_exp_mat_2(samplelist, genelist, pipe.tableName)
    # Create sample index
    expMat.create_sample_index()

    # Normalize the expression values data
    # if quantile normalization or other was used it needs to be included here.
    if norm:
        #expMat.log2_norm()
        if row_aveg:
            expMat.normalize_by_avg_expval()
            
    else:
        # It normalize the expression of each sample according to the expression 
        # level of the normal tissue.
        #expMat.log2_norm()
        #expMat.normalizeexp_bynormals(normal_samples)
        expMat.normalizeexp_bynormals_raw(normal_samples)
        
    return expMat


def write_expmat_file(expMat, matfile, factors, standarization, genespair=[]):
    """
    writes a expression matrix. 
    """
    expMat.create_gen_index()    
    samples = [ i[1] for i in expMat.sample]
    matfile.write('gene'+'\t'+'load'+'\t'+",".join(samples).replace(",", "\t")+'\n')
    #print genespair
    
    for i, g in enumerate(range(expMat.expVal.shape[0])):
        g = expMat.geneSym[i]
        #gind=expMat.genInd[g]
        
        if factors and genespair:
            gload=genespair[i][1]
        else:
            gload='NA'
        
        #print i, gload
        if standarization:
            mat2 = expMat.expVal[i,:]
            # standarize the value of the expression matrix between -1 to 1. 
            el = map( str,list( 2*(mat2/float((np.max(mat2) - np.min(mat2))))) )
            
        else:
            el = map( str,list( expMat.expVal[i,:].filled(fill_value=0)) )
            #el = map( str,list( expMat.expVal[i,:]) )
            
            
        matfile.write(g+'\t'+str(gload)+'\t'+",".join(list( el )).replace(",", "\t") + '\n')
        
    

##########

if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--annotFile", dest="annotFile",
                            help="annotation file for all files to use")
    optionparser.add_option("-n", action="store_true", dest="norm",
                            help="normalize log2 and scale expression vals. Default only log2 and scale by normals. Default universe. Bool")
    optionparser.add_option("-t", action="store_true", dest="factors",
                            help="use the genes produced in the factor analysis")
    optionparser.add_option("-s", action="store_true", dest="standard",
                            help="standardize expression value between -1 and 1.use the genes produced in the factor analysis")
    optionparser.add_option("-v", action="store_true", dest="rowavg",
                            help="normalize the expression values by their median value across the samples")




    
    (options, args) = optionparser.parse_args()
       
    # start the pipeline. 
    # open the database
    pipe = pipeline("localhost", "oabalbin", "oscar", "rnaseq_signatures")
    seed_file, samples, outdirs, thisfactors = pipe.parse_input_annotfiles(open(options.annotFile))    
    #read folder for expression profiles
    
   
    ######### Get raw data  ############
    # Get the database files
    tp = pm.parser()
    pipe.mt.create_connector()
    pipe.mq.create_connector()
    gtp = gpm.geneparser()
    
    
    ########## Make queries ############
    
    # list of samples:
    #samplelist = tp.list_of_names_in_line(open(samples['samples']))
    #normal_samples=tp.list_of_names_in_line(open(samples['normals']))
    samplelist = tp.list_of_names(open(samples['samples']))
    normal_samples=list(tp.list_of_names(open(samples['normals'])))
    
     
    if 'spannot' in samples.keys():
        spannot = tp.list_of_samples_status(open(samples['spannot']))  
        pipe.spclass0=spannot['ETS+']
        pipe.spclass1=spannot['ETS-']
        sampleclasses = [pipe.spclass0, pipe.spclass1]
        samplelist = tl.get_reordered_samplelist(samplelist,sampleclasses)
    
    if options.factors:
        #factorlist.append( bfrm_factor(facnum,genedic,sdgene) )
        factorlist, allgenes = gtp.list_of_nameload_in_factor(open(seed_file['univfile']),thisfactors['usefactors'])    
        controls =['ERG','PLAT','ETV1','ETV4','ETV5','PLA1A','MMP9','SPINK1']
        conrolspairs = [('ERG', 0.0),('PLAT', 0.0),('ETV1', 0.0),('ETV4', 0.0),('ETV5', 0.0),('PLA1A', 0.0),('MMP9', 0.0),('SPINK1', 0.0)]
        
        for ft in factorlist:
            outfile = outdirs['outdir']+'mat_expfactor_'+str(ft.facnum)+'norm_test.txt'
            
            if options.factors:
                genespair = sorted(ft.genedic.items(), key=operator.itemgetter(1), reverse=True)
                genelist = deque([ gp[0] for gp in genespair])
                genelist.extend(controls)
                genespair.extend(conrolspairs)
            
            #print len(genespair)
            expMat = get_expMat(pipe, samplelist, genelist, normal_samples, options.norm, options.rowavg)
            ########### Print expression matrix
            # write expression matrix
            write_expmat_file(expMat, open(outfile,'w'),options.factors, options.standard, genespair)

    else:
        genelist = deque(set(tp.list_of_names(open(seed_file['univfile']))))
        #genelist = tp.list_of_names(open(seed_file['univfile']))
        print genelist
        
        outfile = outdirs['outdir']+'mat_expfactor_'+str('allTranscriptsProstateTissuev17_5_')+'norm_raw_jun8.txt'
        #print len(genespair)
        expMat = get_expMat(pipe, samplelist, genelist, normal_samples, options.norm, options.rowavg)
        ########### Print expression matrix
        # write expression matrix
        write_expmat_file(expMat, open(outfile,'w'),options.factors, options.standard)
        
    sys.exit(0)
